package network

import (
	"adkr/pkg/adkr"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"sync"
	"time"
)

// NetConfig, Network struct, NewNetwork, Start, WaitForFullConnection, listen are unchanged.
// ... (code from previous correct version) ...

// MODIFIED: dispatch now handles loopback, broadcast, and point-to-point messages.
func (n *Network) dispatch() {
	for msg := range n.MsgOutCh {
		msg.Sender = n.config.ID // Set the sender ID before dispatching

		if msg.To == n.config.ID {
			// This is a loopback message, intended for ourself.
			// Put it directly into the input channel.
			log.Printf("Node %d: Loopback message for self (Round %d, Tag %s)", n.config.ID, msg.Round, msg.Tag)
			n.MsgInCh <- msg
		} else if msg.To == -1 {
			// Broadcast message
			n.Broadcast(msg)
		} else {
			// Point-to-point message
			n.Send(msg.To, msg)
		}
	}
}

// ... (The rest of the file is identical to your current working version)
// For clarity, here is the full, final version of the file.
type NetConfig struct {
	ID    int
	Addrs []string
}
type Network struct {
	config   NetConfig
	MsgInCh  chan adkr.Message
	MsgOutCh chan adkr.Message
	conns    map[int]net.Conn
	connMux  sync.RWMutex
}

func NewNetwork(config NetConfig, msgInCh, msgOutCh chan adkr.Message) *Network {
	return &Network{config: config, MsgInCh: msgInCh, MsgOutCh: msgOutCh, conns: make(map[int]net.Conn)}
}
func (n *Network) Start() {
	go n.listen()
	go n.dialPeers()
	go n.dispatch()
	log.Printf("Network layer started for Node %d", n.config.ID)
}
func (n *Network) WaitForFullConnection() {
	log.Printf("Node %d is waiting for a full network connection...", n.config.ID)
	for {
		n.connMux.RLock()
		numConns := len(n.conns)
		n.connMux.RUnlock()
		if numConns >= len(n.config.Addrs)-1 {
			log.Printf("Node %d has a full network connection.", n.config.ID)
			return
		}
		time.Sleep(1 * time.Second)
	}
}
func (n *Network) listen() {
	addr := n.config.Addrs[n.config.ID]
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		log.Fatalf("Node %d failed to listen on %s: %v", n.config.ID, addr, err)
	}
	defer listener.Close()
	log.Printf("Node %d is listening on %s", n.config.ID, addr)
	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Printf("Node %d accept failed: %v", n.config.ID, err)
			continue
		}
		go n.handleConnection(conn)
	}
}
func (n *Network) dialPeers() {
	for peerID, addr := range n.config.Addrs {
		if peerID <= n.config.ID {
			continue
		}
		go func(peerID int, addr string) {
			for {
				n.connMux.RLock()
				_, connected := n.conns[peerID]
				n.connMux.RUnlock()
				if connected {
					time.Sleep(5 * time.Second)
					continue
				}
				conn, err := net.DialTimeout("tcp", addr, 2*time.Second)
				if err != nil {
					time.Sleep(2 * time.Second)
					continue
				}
				n.handleConnection(conn)
			}
		}(peerID, addr)
	}
}
func (n *Network) Send(to int, msg adkr.Message) {
	n.connMux.RLock()
	conn, ok := n.conns[to]
	n.connMux.RUnlock()
	if !ok {
		log.Printf("Node %d: No connection to peer %d, dropping message.", n.config.ID, to)
		return
	}
	if err := n.writeMessage(conn, msg); err != nil {
		log.Printf("Node %d failed to send message to Node %d: %v", n.config.ID, to, err)
	}
}
func (n *Network) Broadcast(msg adkr.Message) {
	n.connMux.RLock()
	connsCopy := make(map[int]net.Conn)
	for id, conn := range n.conns {
		connsCopy[id] = conn
	}
	n.connMux.RUnlock()
	for peerID, conn := range connsCopy {
		if err := n.writeMessage(conn, msg); err != nil {
			log.Printf("Node %d failed to broadcast to Node %d: %v", n.config.ID, peerID, err)
		}
	}
}
func (n *Network) writeMessage(conn net.Conn, msg adkr.Message) error {
	msgBytes, err := json.Marshal(msg)
	if err != nil {
		return fmt.Errorf("failed to marshal message: %w", err)
	}
	lenBuf := make([]byte, 4)
	binary.BigEndian.PutUint32(lenBuf, uint32(len(msgBytes)))
	if _, err := conn.Write(lenBuf); err != nil {
		return fmt.Errorf("failed to write message length: %w", err)
	}
	if _, err := conn.Write(msgBytes); err != nil {
		return fmt.Errorf("failed to write message payload: %w", err)
	}
	return nil
}
func (n *Network) handleConnection(conn net.Conn) {
	if _, err := conn.Write([]byte{byte(n.config.ID)}); err != nil {
		conn.Close()
		return
	}
	idBuf := make([]byte, 1)
	if _, err := io.ReadFull(conn, idBuf); err != nil {
		conn.Close()
		return
	}
	peerID := int(idBuf[0])
	n.connMux.Lock()
	if oldConn, ok := n.conns[peerID]; ok {
		oldConn.Close()
	}
	n.conns[peerID] = conn
	n.connMux.Unlock()
	log.Printf("Node %d established stable connection with Node %d", n.config.ID, peerID)
	defer func() {
		n.connMux.Lock()
		if n.conns[peerID] == conn {
			delete(n.conns, peerID)
		}
		n.connMux.Unlock()
		conn.Close()
		log.Printf("Connection between Node %d and Node %d is closed.", n.config.ID, peerID)
	}()
	for {
		lenBuf := make([]byte, 4)
		if _, err := io.ReadFull(conn, lenBuf); err != nil {
			break
		}
		msgLen := binary.BigEndian.Uint32(lenBuf)
		msgBuf := make([]byte, msgLen)
		if _, err := io.ReadFull(conn, msgBuf); err != nil {
			break
		}
		var msg adkr.Message
		if err := json.Unmarshal(msgBuf, &msg); err != nil {
			log.Printf("Unmarshal error from %d: %v", peerID, err)
			continue
		}
		msg.Sender = peerID
		n.MsgInCh <- msg
	}
}
